﻿using SDILReader;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Threading.Tasks;

namespace DllCallAnalysis45
{
    public static class MethodAnalysisHelper
    {
        /// <summary>
        /// 获取这个函数调用的函数
        /// </summary>
        /// <param name="allTypes"></param>
        /// <param name="method"></param>
        /// <param name="methodFilter"></param>
        /// <returns></returns>
        public static List<MethodInfo> GetMethodCallMethods(Type[] allTypes, MethodInfo method, Func<Assembly, bool> assemblyFilter = null, Func<MethodInfo, bool> methodFilter = null)
        {
            var callInfos = new List<MethodInfo>();

            SDILReader.MethodBodyReader mr = new MethodBodyReader(method);

            for (int i = 0; i < (mr?.instructions?.Count ?? 0); i++)
            {
                var il = mr.instructions[i];

                var doName = il.Code.ToString();

                if (doName == "call")
                {
                    var callMethod = il.Operand as MethodInfo;
                    if (callMethod != null)
                    {
                        if (methodFilter == null || methodFilter(method))
                            callInfos.Add(callMethod);
                    }
                }
                else if (doName == "callvirt")
                {
                    var callMethod = il.Operand as MethodInfo;
                    if (callMethod != null)
                    {
                        if (methodFilter == null || methodFilter(callMethod))
                        {
                            if (callMethod.DeclaringType.IsInterface || callMethod.IsAbstract)
                            {
                                // 尝试检查所有实现的子类
                                var canMatchTypes = allTypes.Where(n => callMethod.DeclaringType.IsAssignableFrom(n)).ToList();
                                if (canMatchTypes?.Any() == true)
                                {
                                    foreach (var matchType in canMatchTypes)
                                    {
                                        if (assemblyFilter != null && assemblyFilter.Invoke(matchType.Assembly) != true)
                                            continue;

                                        if (callMethod.DeclaringType == matchType)
                                            continue;

                                        var matchMethods = matchType.GetMethods()
                                           .Where(n => IsSameMethod(n, callMethod))
                                           .ToList();

                                        foreach (var matchCallMethod in matchMethods)
                                        {
                                            callInfos.Add(callMethod);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            return callInfos;
        }

        /// <summary>
        /// 获取全部的函数调用信息
        /// </summary>
        /// <param name="assemblies"></param>
        /// <param name="mathodFilter"></param>
        /// <returns></returns>
        public static Dictionary<MethodInfo, List<MethodInfo>> GetAllMethodCallInfo(Assembly[] assemblies, Func<Assembly, bool> assemblyFilter, Func<MethodInfo, bool> mathodFilter)
        {
            var mapping = new Dictionary<MethodInfo, List<MethodInfo>>();
            var allTypes = new List<Type>();

            foreach (var assembly in assemblies)
            {
                try
                {
                    if (assemblyFilter != null && assemblyFilter.Invoke(assembly) != true)
                        continue;

                    var types = assembly.GetTypes();
                    if (types?.Any() == true)
                        allTypes.AddRange(types);
                }
                catch (Exception ex)
                {
                    Console.WriteLine($"尝试加载程序集({assembly.FullName})出错:{ex}");
                }
            }

            foreach (var type in allTypes)
            {
                var methods = type.GetMethods();

                foreach (var method in methods)
                {
                    try
                    {
                        var callMethods = GetMethodCallMethods(allTypes.ToArray(), method, assemblyFilter, mathodFilter);
                        if (callMethods?.Any() != true)
                            continue;

                        mapping.Add(method, callMethods);
                    }
                    catch (Exception ex)
                    {
                        Console.WriteLine($"尝试加载函数({method.DeclaringType.FullName}.{method.Name})出错:{ExceptionHelper.BuildExceptionMessage(ex)}");
                    }
                }
            }

            return mapping;
        }

        /// <summary>
        /// 反转调用与被调用的映射关系
        /// </summary>
        /// <param name="callInfos"></param>
        /// <returns></returns>
        public static Dictionary<MethodInfo, List<MethodInfo>> ReversalCallInfo(Dictionary<MethodInfo, List<MethodInfo>> callInfos)
        {
            if (callInfos?.Any() != true)
                return null;

            var calledInfos = new Dictionary<MethodInfo, List<MethodInfo>>();

            foreach (var callInfo in callInfos)
            {
                foreach (var method in callInfo.Value)
                {
                    if (calledInfos.TryGetValue(method, out var calledInfo) != true)
                    {
                        calledInfo = new List<MethodInfo>();
                        calledInfos[method] = calledInfo;
                    }

                    calledInfo.Add(callInfo.Key);
                }
            }

            return calledInfos;
        }

        //public static string BuildMethodName(MethodInfo method)
        //{
        //    return $"{method.ReflectedType}::{method.Name}";
        //}

        public static string BuildCallTree(Type[] allTypes, MethodInfo _method, Func<MethodInfo, bool> filter = null)
        {
            var touchMethods = new List<MethodInfo>();

            var sb = new StringBuilder();

            _BuildCallTree(_method);

            return sb.ToString();

            void _BuildCallTree(MethodInfo method, int deep = 0)
            {
                if (touchMethods.Contains(method))
                    return;

                touchMethods.Add(method);

                SDILReader.MethodBodyReader mr = new MethodBodyReader(method);

                for (int i = 0; i < (mr?.instructions?.Count ?? 0); i++)
                {
                    var il = mr.instructions[i];

                    var doName = il.Code.ToString();

                    if (doName == "call")
                    {
                        var callMethod = il.Operand as MethodInfo;
                        if (IsPassFilter(callMethod))
                        {
                            sb.AppendLine($"{GetDeepHead(deep)}{il.GetCode()}");

                            _BuildCallTree(callMethod, deep + 1);
                        }
                    }
                    else if (doName == "callvirt")
                    {
                        var callMethod = il.Operand as MethodInfo;
                        if (callMethod != null)
                        {
                            if (IsPassFilter(callMethod))
                            {
                                sb.AppendLine($"{GetDeepHead(deep)}{il.GetCode()}");
                                _BuildCallTree(callMethod, deep + 1);

                                if (callMethod.DeclaringType.IsInterface || callMethod.IsAbstract)
                                {
                                    // 尝试检查所有实现的子类
                                    var canMatchTypes = allTypes.Where(n => callMethod.DeclaringType.IsAssignableFrom(n)).ToList();
                                    if (canMatchTypes?.Any() == true)
                                    {
                                        foreach (var matchType in canMatchTypes)
                                        {
                                            if (callMethod.DeclaringType == matchType)
                                                continue;

                                            var matchMethods = matchType.GetMethods()
                                                .Where(n => IsPassFilter(n))
                                                .Where(n => IsSameMethod(n, callMethod))
                                                .ToList();

                                            foreach (var matchCallMethod in matchMethods)
                                            {
                                                sb.AppendLine($"{GetDeepHead(deep + 1)}calltrue {BuildMethodCallCode(matchCallMethod)}");
                                                _BuildCallTree(matchCallMethod, deep + 2);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            bool IsPassFilter(MethodInfo method)
            {
                if (filter == null)
                    return true;

                if (method == null)
                    return false;

                return filter.Invoke(method);
            }
        }

        public static string TrimCallTree(string tree, params string[] matchText)
        {
            var lines = tree.Split(new char[] { '\r', '\n' }, StringSplitOptions.RemoveEmptyEntries);


            DeepCode deepCode = new DeepCode { Deep = -1 };

            foreach (var line in lines)
            {
                var lDeep = GetDeep(line);

                var _deepCode = FindOrCreateDeepCode(deepCode, lDeep);

                _deepCode.Lines.Add(line);
            }

            var sb = new StringBuilder();

            _Trim(deepCode);

            return sb.ToString();

            void _Trim(DeepCode _deepCode)
            {
                if (_deepCode.Lines.Any())
                {
                    if (IsCurrentNotTrim(_deepCode))
                    {
                        foreach (var line in _deepCode.Lines)
                        {
                            if (IsNotTrimText(line))
                                sb.AppendLine(line);
                        }
                    }
                    else if (IsNotTrim(_deepCode))
                    {
                        sb.AppendLine(_deepCode.Lines.Last());
                    }
                }

                foreach (var deep in _deepCode.DeepCodes)
                {
                    _Trim(deep);
                }
            }

            bool IsNotTrimText(string text)
            {
                return matchText.Any(n => text.Contains(n));
            }

            bool IsCurrentNotTrim(DeepCode _deepCode)
            {
                foreach (var item in _deepCode.Lines)
                {
                    if (IsNotTrimText(item))
                        return true;
                }

                return false;
            }

            bool IsNotTrim(DeepCode _deepCode)
            {
                if (IsCurrentNotTrim(_deepCode))
                    return true;

                foreach (var item in _deepCode.DeepCodes)
                {
                    if (IsNotTrim(item))
                        return true;
                }

                return false;
            }
        }

        public static DeepCode FindOrCreateDeepCode(DeepCode deepCode, int deep)
        {
            return _FindOrCreateDeepCode(deepCode);

            DeepCode _FindOrCreateDeepCode(DeepCode _deepCode)
            {
                // 到达当前层级的上级了
                if (_deepCode.Deep == deep - 1)
                {
                    if (_deepCode.DeepCodes?.Any() == true)
                    {
                        var lastChild = _deepCode.DeepCodes.Last();
                        if (lastChild.DeepCodes.Any())
                        {
                            var newDeepCode = new DeepCode { Deep = _deepCode.Deep + 1 };

                            _deepCode.DeepCodes.Add(newDeepCode);

                            return newDeepCode;
                        }
                    }
                }

                // 当前层级就是了
                if (_deepCode.Deep == deep)
                {
                    return _deepCode;
                }

                if (_deepCode.DeepCodes.Any() != true)
                {
                    var newDeepCode = new DeepCode { Deep = _deepCode.Deep + 1 };

                    _deepCode.DeepCodes.Add(newDeepCode);
                }

                return _FindOrCreateDeepCode(_deepCode.DeepCodes.Last());
            }
        }

        public static int GetDeep(string line)
        {
            var count = 0;

            for (int i = 0; i < line.Length; i++)
            {
                if (line[i] != ' ')
                {
                    count = i;
                    break;
                }
            }

            return count / 2;
        }

        public static string BuildMethodCallCode(MethodInfo mi)
        {
            return Globals.ProcessSpecialTypes(mi.ReturnType.ToString()) + " " + Globals.ProcessSpecialTypes(mi.ReflectedType.ToString()) + "::" + mi.Name + "()";
        }

        public static string GetDeepHead(int deep)
        {
            var sb = new StringBuilder();

            for (int i = 0; i < deep; i++)
            {
                sb.Append("  ");
            }

            return sb.ToString();
        }

        public static bool IsSameMethod(MethodInfo a, MethodInfo b)
        {
            if (a == null || b == null)
                return false;

            if (a.Name != b.Name)
                return false;

            var aParameters = a.GetParameters();
            var bParameters = b.GetParameters();

            if (aParameters.Length != bParameters.Length)
                return false;

            for (int i = 0; i < aParameters.Length; i++)
            {
                if (aParameters[i].ParameterType != bParameters[i].ParameterType)
                    return false;
            }

            if (a.ReturnType != b.ReturnType)
                return false;

            return true;
        }
    }

    public class DeepCode
    {
        public int Deep { get; set; }

        public List<string> Lines { get; set; } = new List<string>();

        public List<DeepCode> DeepCodes { get; set; } = new List<DeepCode>();
    }
}
